package com.janetfilter.core;

import com.janetfilter.core.commons.DebugInfo;
import com.janetfilter.core.plugin.MyTransformer;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.*;

/**
 * 分发器类，实现ClassFileTransformer接口，是整个框架的核心
 * 负责管理所有转换器并将转换请求分发给相应的转换器
 * 同时也负责拦截native方法调用（通过实现ClassFileTransformer接口的同名方法）
 */
public final class Dispatcher implements ClassFileTransformer {
    // 环境配置
    private final Environment environment;
    // 记录所有需要转换的类名
    private final Set<String> classSet = new TreeSet<>();
    // 记录针对特定类的转换器映射，key是类名，value是转换器列表
    private final Map<String, List<MyTransformer>> transformerMap = new HashMap<>();
    // 全局转换器列表，这些转换器会处理所有类
    private final List<MyTransformer> globalTransformers = new ArrayList<>();
    // 管理型转换器列表，这些转换器会管理转换过程
    private final List<MyTransformer> manageTransformers = new ArrayList<>();

    /**
     * 创建分发器实例
     * 
     * @param environment 环境配置
     */
    public Dispatcher(Environment environment) {
        this.environment = environment;
    }

    /**
     * 添加一个转换器
     * 
     * @param transformer 要添加的转换器
     */
    public void addTransformer(MyTransformer transformer) {
        if (null == transformer) {
            return;
        }

        // 检查转换器是否支持当前模式（Attach或Java Agent）
        if (environment.isAttachMode() && !transformer.attachMode()) {
            DebugInfo.debug("Transformer: " + transformer.getClass().getName() + " is set to not load in attach mode, ignored.");
            return;
        }

        if (environment.isJavaagentMode() && !transformer.javaagentMode()) {
            DebugInfo.debug("Transformer: " + transformer.getClass().getName() + " is set to not load in -javaagent mode, ignored.");
            return;
        }

        synchronized (this) {
            // 获取转换器要处理的类名
            String className = transformer.getHookClassName();
            if (null == className) {
                // 如果类名为null，表示这是一个全局转换器
                globalTransformers.add(transformer);

                // 如果是管理型转换器，添加到管理型转换器列表
                if (transformer.isManager()) {
                    manageTransformers.add(transformer);
                }

                return;
            }

            // 将类名添加到需要转换的类集合中（将/替换为.，使其符合Java的类名格式）
            classSet.add(className.replace('/', '.'));
            // 将转换器添加到对应类名的转换器列表中
            List<MyTransformer> transformers = transformerMap.computeIfAbsent(className, k -> new ArrayList<>());

            transformers.add(transformer);
        }
    }

    /**
     * 添加多个转换器
     * 
     * @param transformers 转换器列表
     */
    public void addTransformers(List<MyTransformer> transformers) {
        if (null == transformers) {
            return;
        }

        for (MyTransformer transformer : transformers) {
            addTransformer(transformer);
        }
    }

    /**
     * 添加多个转换器
     * 
     * @param transformers 转换器数组
     */
    public void addTransformers(MyTransformer[] transformers) {
        if (null == transformers) {
            return;
        }

        addTransformers(Arrays.asList(transformers));
    }

    /**
     * 获取所有需要转换的类名
     * 
     * @return 类名集合
     */
    public Set<String> getHookClassNames() {
        return classSet;
    }

    /**
     * 实现ClassFileTransformer接口的transform方法
     * 当类被加载时，JVM会调用此方法进行字节码转换
     * 
     * @param loader 类加载器
     * @param className 类名
     * @param classBeingRedefined 被重定义的类
     * @param protectionDomain 保护域
     * @param classFileBuffer 原始字节码
     * @return 转换后的字节码
     * @throws IllegalClassFormatException 如果字节码格式不合法
     */
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classFileBuffer) throws IllegalClassFormatException {
        if (null == className) {
            return classFileBuffer;
        }

        // 获取针对当前类的特定转换器
        List<MyTransformer> transformers = transformerMap.get(className);
        // 确定要使用的全局转换器列表
        List<MyTransformer> globalTransformers = null == transformers ? this.manageTransformers : this.globalTransformers;

        // 记录转换顺序
        int order = 0;

        try {
            // 1. 调用所有全局转换器的before方法
            for (MyTransformer transformer : globalTransformers) {
                transformer.before(loader, classBeingRedefined, protectionDomain, className, classFileBuffer);
            }

            // 2. 调用所有全局转换器的preTransform方法
            for (MyTransformer transformer : globalTransformers) {
                classFileBuffer = transformer.preTransform(loader, classBeingRedefined, protectionDomain, className, classFileBuffer, order++);
            }

            // 3. 调用针对特定类的转换器
            if (null != transformers) {
                for (MyTransformer transformer : transformers) {
                    classFileBuffer = transformer.transform(loader, classBeingRedefined, protectionDomain, className, classFileBuffer, order++);
                }
            }

            // 4. 调用所有全局转换器的postTransform方法
            for (MyTransformer transformer : globalTransformers) {
                classFileBuffer = transformer.postTransform(loader, classBeingRedefined, protectionDomain, className, classFileBuffer, order++);
            }

            // 5. 调用所有全局转换器的after方法
            for (MyTransformer transformer : globalTransformers) {
                transformer.after(loader, classBeingRedefined, protectionDomain, className, classFileBuffer);
            }
        } catch (Throwable e) {
            DebugInfo.error("Transform class failed: " + className, e);
        }

        // 返回最终转换后的字节码
        return classFileBuffer;
    }
}
